import os
import cv2
import numpy as np
from pathlib import Path
import argparse

def create_black_images(src_dir, dest_dir):
    """
    为源目录中的每张图片创建相同尺寸的全黑图片，并保存到目标目录
    
    Args:
        src_dir (str): 源图片目录路径
        dest_dir (str): 目标目录路径
    """
    # 确保目标目录存在
    Path(dest_dir).mkdir(parents=True, exist_ok=True)
    
    # 支持的图片格式
    img_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']
    
    # 遍历源目录中的所有文件
    for filename in os.listdir(src_dir):
        # 检查文件扩展名是否为图片格式
        if any(filename.lower().endswith(ext) for ext in img_extensions):
            # 构建源文件路径
            src_path = os.path.join(src_dir, filename)
            
            # 读取图片以获取尺寸
            img = cv2.imread(src_path)
            if img is None:
                print(f"无法读取图片: {src_path}")
                continue
                
            # 获取图片尺寸
            height, width = img.shape[:2]
            
            # 创建全黑图片
            black_img = np.zeros((height, width, 3), dtype=np.uint8)
            
            # 构建目标文件路径
            dest_path = os.path.join(dest_dir, filename)
            
            # 保存全黑图片
            cv2.imwrite(dest_path, black_img)
            print(f"已创建: {dest_path} (尺寸: {width}x{height})")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="为源目录中的每张图片创建相同尺寸的全黑图片")
    parser.add_argument("src_dir", help="源图片目录路径")
    parser.add_argument("dest_dir", help="目标目录路径")
    
    args = parser.parse_args()
    
    create_black_images(args.src_dir, args.dest_dir)
    print("完成！")